//
// rs_color_depth.c
//
// License: Apache 2.0. See LICENSE file in root directory.
//
// This example shows how to capture both COLOR and DEPTH frames in the
// same application using a single pipeline, and write each new frame out to
// its own binary file (color_data.bin, depth_data.bin). Every frame update
// completely overwrites the previous file, so each file always contains only
// the latest frame.
//

#include <librealsense2/rs.h>
#include <librealsense2/h/rs_pipeline.h>
#include <librealsense2/h/rs_frame.h>
#include <librealsense2/h/rs_option.h>
#include <stddef.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include "example.h"

//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
// Configuration Macros
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#define DEPTH_STREAM           RS2_STREAM_DEPTH   
#define DEPTH_FORMAT           RS2_FORMAT_Z16     
#define DEPTH_WIDTH            640               
#define DEPTH_HEIGHT           480               
#define DEPTH_FPS              30                
#define DEPTH_STREAM_INDEX     0                 

#define COLOR_STREAM           RS2_STREAM_COLOR 
#define COLOR_FORMAT           RS2_FORMAT_RGB8  
#define COLOR_WIDTH            640             
#define COLOR_HEIGHT           480             
#define COLOR_FPS              30              
#define COLOR_STREAM_INDEX     0               

// Enable or disable writing out data
#define SAVE_DEPTH_DATA        1
#define SAVE_COLOR_DATA        1

//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
// Some data headers to identify what we've written to the bin files
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
typedef struct {
    uint32_t magic;            // Magic number for file identification
    uint32_t width;            // Frame width
    uint32_t height;           // Frame height
    uint32_t format;           // Data format (RS2_FORMAT_Z16 for depth)
    uint32_t bytes_per_pixel;  // 2 for Z16 depth
    uint32_t frame_count;      // Number of frames in file (here, 1)
} depth_file_header_t;

typedef struct {
    uint32_t magic;            // Magic number for file identification
    uint32_t width;            // Frame width
    uint32_t height;           // Frame height
    uint32_t format;           // Data format (RS2_FORMAT_RGB8 for color)
    uint32_t bytes_per_pixel;  // 3 for RGB8
    uint32_t frame_count;      // Number of frames in file (here, 1)
} color_file_header_t;

//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
// Utility: get depth scale (depth_units)
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
static float get_depth_unit_value(const rs2_device* const dev)
{
    rs2_error* e = 0;
    rs2_sensor_list* sensor_list = rs2_query_sensors(dev, &e);
    check_error(e);

    int num_of_sensors = rs2_get_sensors_count(sensor_list, &e);
    check_error(e);

    float depth_scale = 0.0f;
    int is_depth_sensor_found = 0;

    for (int i = 0; i < num_of_sensors; ++i)
    {
        rs2_sensor* sensor = rs2_create_sensor(sensor_list, i, &e);
        check_error(e);

        is_depth_sensor_found = rs2_is_sensor_extendable_to(sensor, RS2_EXTENSION_DEPTH_SENSOR, &e);
        check_error(e);

        if (is_depth_sensor_found == 1)
        {
            depth_scale = rs2_get_option((const rs2_options*)sensor, RS2_OPTION_DEPTH_UNITS, &e);
            check_error(e);
            rs2_delete_sensor(sensor);
            break;
        }
        rs2_delete_sensor(sensor);
    }

    rs2_delete_sensor_list(sensor_list);

    if (!is_depth_sensor_found)
    {
        printf("Depth sensor not found!\n");
        exit(EXIT_FAILURE);
    }

    return depth_scale;
}

//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
// Main
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
int main()
{
    rs2_error* e = 0;

    // Create a context
    rs2_context* ctx = rs2_create_context(RS2_API_VERSION, &e);
    check_error(e);

    // Get device list
    rs2_device_list* device_list = rs2_query_devices(ctx, &e);
    check_error(e);

    int dev_count = rs2_get_device_count(device_list, &e);
    check_error(e);

    printf("There are %d connected RealSense devices.\n", dev_count);
    if (dev_count == 0)
    {
        rs2_delete_device_list(device_list);
        rs2_delete_context(ctx);
        return EXIT_FAILURE;
    }

    // Use the first detected device
    rs2_device* dev = rs2_create_device(device_list, 0, &e);
    check_error(e);

    // Print some info about the device
    print_device_info(dev);

    // For completeness, retrieve the depth scale
    float depth_scale = get_depth_unit_value(dev);
    printf("Depth scale is %f meters per 1 unit.\n", depth_scale);

    // Create a pipeline
    rs2_pipeline* pipeline = rs2_create_pipeline(ctx, &e);
    check_error(e);

    // Create a config and request both depth & color streams
    rs2_config* config = rs2_create_config(&e);
    check_error(e);

    // Enable depth
    rs2_config_enable_stream(config, 
                             DEPTH_STREAM, 
                             DEPTH_STREAM_INDEX, 
                             DEPTH_WIDTH, 
                             DEPTH_HEIGHT, 
                             DEPTH_FORMAT, 
                             DEPTH_FPS, 
                             &e);
    check_error(e);

    // Enable color
    rs2_config_enable_stream(config, 
                             COLOR_STREAM, 
                             COLOR_STREAM_INDEX, 
                             COLOR_WIDTH, 
                             COLOR_HEIGHT, 
                             COLOR_FORMAT, 
                             COLOR_FPS, 
                             &e);
    check_error(e);

    // Start streaming
    rs2_pipeline_profile* pipeline_profile = rs2_pipeline_start_with_config(pipeline, config, &e);
    if (e)
    {
        printf("The connected device doesn't support the requested streams!\n");
        exit(EXIT_FAILURE);
    }

    printf("Streaming color + depth frames...\n");
    printf("Each new depth frame overwrites depth_data.bin, and each new color frame overwrites color_data.bin.\n");

    // Start looping for frames
    while (1)
    {
        rs2_frame* frames = rs2_pipeline_wait_for_frames(pipeline, RS2_DEFAULT_TIMEOUT, &e);
        check_error(e);

        int num_of_frames = rs2_embedded_frames_count(frames, &e);
        check_error(e);

        for (int i = 0; i < num_of_frames; ++i)
        {
            rs2_frame* frame = rs2_extract_frame(frames, i, &e);
            check_error(e);

            // Identify the stream type
            const rs2_stream_profile* sp = rs2_get_frame_stream_profile(frame, &e);
            check_error(e);

            rs2_stream stream_type;
            rs2_format fmt;
            int idx, uid, fr;
            rs2_get_stream_profile_data(sp, &stream_type, &fmt, &idx, &uid, &fr, &e);
            check_error(e);

            if (stream_type == RS2_STREAM_DEPTH && SAVE_DEPTH_DATA)
            {
                // Depth frame
                const uint16_t* depth_data = (const uint16_t*) rs2_get_frame_data(frame, &e);
                check_error(e);

                int depth_w = rs2_get_frame_width(frame, &e);
                check_error(e);
                int depth_h = rs2_get_frame_height(frame, &e);
                check_error(e);

                // Write to .tmp file, then rename
                FILE* depth_file_tmp = fopen("depth_data.bin.tmp", "wb");
                if (depth_file_tmp)
                {
                    depth_file_header_t header;
                    header.magic           = 0x44455054; // 'DEPT'
                    header.width           = depth_w;
                    header.height          = depth_h;
                    header.format          = DEPTH_FORMAT;
                    header.bytes_per_pixel = 2; // Z16
                    header.frame_count     = 1;

                    fwrite(&header, sizeof(header), 1, depth_file_tmp);
                    fwrite(depth_data, sizeof(uint16_t), depth_w * depth_h, depth_file_tmp);
                    fclose(depth_file_tmp);

                    // Atomically rename to final .bin
                    if (rename("depth_data.bin.tmp", "depth_data.bin") != 0)
                    {
                        perror("rename depth_data.bin.tmp -> depth_data.bin");
                    }

                    unsigned long long frame_number = rs2_get_frame_number(frame, &e);
                    check_error(e);
                    printf("Depth frame #%llu -> depth_data.bin with tmp (%dx%d)\n", 
                           frame_number, depth_w, depth_h);
                }
            }
            else if (stream_type == RS2_STREAM_COLOR && SAVE_COLOR_DATA)
            {
                // Color frame
                const uint8_t* color_data = (const uint8_t*) rs2_get_frame_data(frame, &e);
                check_error(e);

                int color_w = rs2_get_frame_width(frame, &e);
                check_error(e);
                int color_h = rs2_get_frame_height(frame, &e);
                check_error(e);

                int stride = rs2_get_frame_stride_in_bytes(frame, &e);
                check_error(e);

                // Write to .tmp file, then rename
                FILE* color_file_tmp = fopen("color_data.bin.tmp", "wb");
                if (color_file_tmp)
                {
                    color_file_header_t header;
                    header.magic           = 0x434F4C52; // 'COLR'
                    header.width           = color_w;
                    header.height          = color_h;
                    header.format          = COLOR_FORMAT;
                    header.bytes_per_pixel = 3; // RGB8
                    header.frame_count     = 1;

                    fwrite(&header, sizeof(header), 1, color_file_tmp);

                    // If the stride is exactly width*3, we can write row by row quickly
                    if (stride == color_w * 3)
                    {
                        fwrite(color_data, sizeof(uint8_t), color_w * color_h * 3, color_file_tmp);
                    }
                    else
                    {
                        // Otherwise, handle row by row
                        for (int row = 0; row < color_h; row++)
                        {
                            fwrite(color_data + row * stride, sizeof(uint8_t), color_w * 3, color_file_tmp);
                        }
                    }
                    fclose(color_file_tmp);

                    // Atomically rename to final .bin
                    if (rename("color_data.bin.tmp", "color_data.bin") != 0)
                    {
                        perror("rename color_data.bin.tmp -> color_data.bin");
                    }

                    unsigned long long frame_number = rs2_get_frame_number(frame, &e);
                    check_error(e);
                    printf("Color frame #%llu -> color_data.bin with tmp (%dx%d)\n", 
                           frame_number, color_w, color_h);
                }
            }

            rs2_release_frame(frame);
        }

        rs2_release_frame(frames);
    }

    // Cleanup
    rs2_pipeline_stop(pipeline, &e);
    check_error(e);

    rs2_delete_pipeline_profile(pipeline_profile);
    rs2_delete_config(config);
    rs2_delete_pipeline(pipeline);
    rs2_delete_device(dev);
    rs2_delete_device_list(device_list);
    rs2_delete_context(ctx);

    return EXIT_SUCCESS;
} 
